{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-01-13 11:42:25.832558: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "E0000 00:00:1736739746.009490    4944 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "E0000 00:00:1736739746.051817    4944 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-01-13 11:42:26.444646: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0049,  0.0576, -0.0072,  ..., -0.0471,  0.0080,  0.0103],\n",
      "        [-0.0073,  0.0228, -0.0478,  ..., -0.0503,  0.0201,  0.0724],\n",
      "        [ 0.0077,  0.0220, -0.0315,  ..., -0.0360,  0.0306,  0.0844]],\n",
      "       device='cuda:0')\n",
      "204739 204739\n"
     ]
    }
   ],
   "source": [
    "# 生成all标签嵌入向量\n",
    "import os\n",
    "import json\n",
    "\n",
    "root='/media/jh/新加卷'\n",
    "with open(f'{root}/2024_12_04/llms4subjects-main/shared-task-datasets/GND/dataset/GND-Subjects-all.json') as f:\n",
    "    gnd_jsons=json.load(f)\n",
    "labels=[]\n",
    "for c in gnd_jsons:\n",
    "    item1={}\n",
    "    item1['sentence2']='Classification Name is '+c['Classification Name']+'. Name is '+c['Name']\n",
    "    item1['Code']=c['Code']\n",
    "    labels.append(item1)\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "model = SentenceTransformer(f'{root}/lastest_out')\n",
    "model.to('cuda')\n",
    "# 把labels的一维张量转换为二维张量\n",
    "batch_size=32\n",
    "texts=[label['sentence2'] for label in labels]\n",
    "texts=[texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]\n",
    "if len(texts)*batch_size<len(labels):\n",
    "    texts.append([label['sentence2'] for label in labels[len(texts)*batch_size:]])\n",
    "y_test=[]\n",
    "for text in texts:\n",
    "    with torch.no_grad():\n",
    "        embeddings = model.encode(text, convert_to_tensor=True)\n",
    "    item=embeddings.cpu()\n",
    "    y_test.append(item)\n",
    "print(embeddings)\n",
    "\n",
    "y_t=[]\n",
    "i=0\n",
    "for y in y_test:\n",
    "    for e in y:\n",
    "        item={}\n",
    "        item['Code']=labels[i]['Code']\n",
    "        item['embedding']=e\n",
    "        item['sentence2']=labels[i]['sentence2']\n",
    "        y_t.append(item)\n",
    "        i+=1\n",
    "print(len(y_t),len(labels)) # 204739 204739\n",
    "\n",
    "import pickle\n",
    "with open(f'{root}/2024_12_04/llms4subjects-main/shared-task-datasets/all_labels.pkl','wb') as f:\n",
    "    pickle.dump(y_t,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0040, -0.0129, -0.0552,  ...,  0.0315,  0.0206,  0.0566],\n",
      "        [-0.0422, -0.0048, -0.0326,  ...,  0.0453, -0.0018,  0.0360],\n",
      "        [-0.0587, -0.0166, -0.0399,  ...,  0.0305, -0.0045,  0.0369]],\n",
      "       device='cuda:0')\n",
      "79427 79427\n"
     ]
    }
   ],
   "source": [
    "# 生成core标签嵌入向量\n",
    "import os\n",
    "import json\n",
    "\n",
    "root='/media/jh/新加卷'\n",
    "with open(f'{root}/2024_12_04/llms4subjects-main/shared-task-datasets/GND/dataset/GND-Subjects-tib-core.json') as f:\n",
    "    gnd_jsons=json.load(f)\n",
    "labels=[]\n",
    "for c in gnd_jsons:\n",
    "    item1={}\n",
    "    item1['sentence2']='Classification Name is '+c['Classification Name']+'. Name is '+c['Name']\n",
    "    item1['Code']=c['Code']\n",
    "    labels.append(item1)\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "model = SentenceTransformer(f'{root}/lastest_out')\n",
    "model.to('cuda')\n",
    "# 把labels的一维张量转换为二维张量\n",
    "batch_size=32\n",
    "texts=[label['sentence2'] for label in labels]\n",
    "texts=[texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]\n",
    "if len(texts)*batch_size<len(labels):\n",
    "    texts.append([label['sentence2'] for label in labels[len(texts)*batch_size:]])\n",
    "y_test=[]\n",
    "for text in texts:\n",
    "    with torch.no_grad():\n",
    "        embeddings = model.encode(text, convert_to_tensor=True)\n",
    "    item=embeddings.cpu()\n",
    "    y_test.append(item)\n",
    "print(embeddings)\n",
    "\n",
    "y_t=[]\n",
    "i=0\n",
    "for y in y_test:\n",
    "    for e in y:\n",
    "        item={}\n",
    "        item['Code']=labels[i]['Code']\n",
    "        item['embedding']=e\n",
    "        item['sentence2']=labels[i]['sentence2']\n",
    "        y_t.append(item)\n",
    "        i+=1\n",
    "print(len(y_t),len(labels)) # 204739 204739\n",
    "\n",
    "import pickle\n",
    "with open(f'{root}/2024_12_04/llms4subjects-main/shared-task-datasets/core_labels.pkl','wb') as f:\n",
    "    pickle.dump(y_t,f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "import os\n",
    "import json\n",
    "import pickle\n",
    "root1='/media/jh/新加卷'\n",
    "with open(f'{root1}/2024_12_04/llms4subjects-main/shared-task-datasets/core_labels.pkl','rb') as f:\n",
    "    y_t=pickle.load(f)\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "model = SentenceTransformer(f'{root1}/lastest_out')\n",
    "model.to('cuda')\n",
    "\n",
    "en_de='en'\n",
    "root='./test-tib-all'\n",
    "\n",
    "des_root=f'{root1}/2024_12_04/result'\n",
    "\n",
    "def get_train_data(json_file):\n",
    "    with open(json_file,'r') as f:\n",
    "        data=json.loads(f.read())\n",
    "    for d in data['@graph']:\n",
    "        if 'title' in d:\n",
    "            tmp=d\n",
    "            break\n",
    "    item={}\n",
    "    if type(tmp['title'])==list:\n",
    "        item['title']=' '.join(tmp['title'])\n",
    "    else:\n",
    "        item['title']=tmp['title']\n",
    "    if type(tmp['abstract'])==list:\n",
    "        item['abstract']='  '.join(tmp['abstract'])\n",
    "    else:\n",
    "        item['abstract']=tmp['abstract']\n",
    "    return item\n",
    "\n",
    "for dirs in os.listdir(f'{root}/'):\n",
    "    for file in os.listdir(f'{root}/{dirs}/{en_de}'):\n",
    "        item=get_train_data(f'{root}/{dirs}/{en_de}/{file}')\n",
    "        text=item['title']+'.'+item['abstract']\n",
    "        with torch.no_grad():\n",
    "            embeddings = model.encode(text, convert_to_tensor=True)\n",
    "        y_tt=torch.tensor([y['embedding'].tolist() for y in y_t]).to('cuda')\n",
    "        with torch.no_grad():\n",
    "            cosine_similarity = F.cosine_similarity(embeddings, y_tt, dim=1)\n",
    "        top_k = torch.topk(cosine_similarity, k=50, dim=0) # 取出相似度最高的50个\n",
    "        top_k=top_k.indices.cpu().tolist()\n",
    "        code=[]\n",
    "        for i in top_k:\n",
    "            code.append(y_t[i]['Code'])\n",
    "        item1={}\n",
    "        item1['dcterms_subject']=code\n",
    "        # 如果不存在目录，则创建目录\n",
    "        if not os.path.exists(f'{des_root}/{dirs}/{en_de}'):\n",
    "            os.makedirs(f'{des_root}/{dirs}/{en_de}')\n",
    "        with open(f'{des_root}/{dirs}/{en_de}/{file}','w') as f:\n",
    "            f.write(json.dumps(item1, indent=4))\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import json\n",
    "# with open('./te.json', 'w') as f:\n",
    "#     code=[1,2,3] \n",
    "#     item1={}\n",
    "#     item1['dcterms_subject']=code\n",
    "#     f.write(json.dumps(item1, indent=4))\n",
    "import os\n",
    "dirs='./result/Book/en'\n",
    "if not os.path.exists(dirs):\n",
    "    os.makedirs(dirs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 制作数据集\n",
    "import pickle\n",
    "import json\n",
    "import random\n",
    "with open(\"all-data/dev_en.pkl\", \"rb\") as f:\n",
    "    test = pickle.load(f)\n",
    "with open(\"all-data/dev_de.pkl\", \"rb\") as f:\n",
    "    test.extend(pickle.load(f))\n",
    "    \n",
    "with open('./GND/dataset/GND-Subjects-all.json') as f:\n",
    "    gnd_jsons=json.load(f)\n",
    "\n",
    "labels=[]\n",
    "for c in gnd_jsons:\n",
    "    item1={}\n",
    "    item1['sentence2']='Classification Name is '+c['Classification Name']+'. Name is '+c['Name']\n",
    "    item1['Code']=c['Code']\n",
    "    labels.append(item1)\n",
    "x_test=[]\n",
    "for t in test:\n",
    "    if type(t['title'])!=str:\n",
    "        t['title']=' '.join(t['title'])\n",
    "        # break\n",
    "    if type(t['abstract'])!=str:\n",
    "        t['abstract']=' '.join(t['abstract'])\n",
    "    item=t['title']+'.'+t['abstract']\n",
    "    x_test.append(item)\n",
    "  \n",
    "from sentence_transformers import SentenceTransformer\n",
    "import torch\n",
    "model = SentenceTransformer('/media/4t/lastest_out')\n",
    "model.to('cuda')\n",
    "# 把labels的一维张量转换为二维张量\n",
    "batch_size=32\n",
    "texts=[label['sentence2'] for label in labels]\n",
    "texts=[texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]\n",
    "if len(texts)*batch_size<len(labels):\n",
    "    texts.append([label['sentence2'] for label in labels[len(texts)*batch_size:]])\n",
    "y_test=[]\n",
    "for text in texts:\n",
    "    with torch.no_grad():\n",
    "        embeddings = model.encode(text, convert_to_tensor=True)\n",
    "    item=embeddings.cpu()\n",
    "    y_test.append(item)\n",
    "print(embeddings)\n",
    "\n",
    "y_t=[]\n",
    "i=0\n",
    "for y in y_test:\n",
    "    for e in y:\n",
    "        item={}\n",
    "        item['Code']=labels[i]['Code']\n",
    "        item['embedding']=e\n",
    "        item['sentence2']=labels[i]['sentence2']\n",
    "        y_t.append(item)\n",
    "        i+=1\n",
    "print(len(y_t),len(labels)) # 204739 204739\n",
    "x_texts=[x_test[i:i+batch_size] for i in range(0, len(x_test), batch_size)]\n",
    "if len(x_texts)*batch_size<len(x_test):\n",
    "    x_texts.append([x_test[len(x_texts)*batch_size:]])\n",
    "x_t=[]\n",
    "for x in x_texts:\n",
    "    with torch.no_grad():\n",
    "        embeddings = model.encode(x, convert_to_tensor=True)\n",
    "    item=embeddings.cpu()\n",
    "    x_t.append(item)\n",
    "print(len(x_t),len(x_texts)) # 6340 6340\n",
    "print(len(x_test))\n",
    "\n",
    "y_tt=torch.tensor([y['embedding'].tolist() for y in y_t]).to('cuda')\n",
    "import torch.nn.functional as F\n",
    "x_scores=[]\n",
    "for x in x_t:\n",
    "    x=x.to('cuda')\n",
    "    for xx in x:\n",
    "        with torch.no_grad():\n",
    "            cosine_similarity = F.cosine_similarity(xx, y_tt, dim=1)\n",
    "        x_scores.append(cosine_similarity.cpu())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
